import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 游戏常量
WINDOW_SIZE = 800
GRID_SIZE = 40
CELL_SIZE = WINDOW_SIZE // GRID_SIZE
GRID_COLOR = (220, 220, 220)  # 浅灰色
SNAKE_COLOR = (0, 255, 0)      # 绿色
SNAKE_HEAD_COLOR = (0, 160, 0) # 深绿色
FOOD_COLOR = (255, 0, 0)       # 红色
TEXT_COLOR = (0, 0, 0)         # 黑色

# 创建游戏窗口
screen = pygame.display.set_mode((WINDOW_SIZE, WINDOW_SIZE))
pygame.display.set_caption("贪吃蛇游戏")
clock = pygame.time.Clock()

# 游戏状态
class GameState:
    def __init__(self):
        self.reset()
    
    def reset(self):
        # 蛇的初始位置（中心点）
        center = GRID_SIZE // 2
        self.snake = [(center, center)]
        self.direction = (1, 0)  # 初始向右移动
        self.next_direction = (1, 0)
        self.food = None
        self.score = 1
        self.game_over = False
        self.generate_food()
    
    def generate_food(self):
        # 生成食物在随机空白位置
        while True:
            food = (random.randint(0, GRID_SIZE-1), random.randint(0, GRID_SIZE-1))
            if food not in self.snake:
                self.food = food
                break
    
    def update(self):
        if self.game_over:
            return
        
        # 更新方向
        self.direction = self.next_direction
        
        # 计算新的蛇头位置
        head_x, head_y = self.snake[0]
        dx, dy = self.direction
        new_head = ((head_x + dx) % GRID_SIZE, (head_y + dy) % GRID_SIZE)
        
        # 检查游戏结束条件
        if new_head in self.snake[1:] or new_head[0] < 0 or new_head[0] >= GRID_SIZE or new_head[1] < 0 or new_head[1] >= GRID_SIZE:
            self.game_over = True
            return
        
        # 移动蛇
        self.snake.insert(0, new_head)
        
        # 检查是否吃到食物
        if new_head == self.food:
            self.score += 1
            self.generate_food()
        else:
            self.snake.pop()

# 创建游戏状态
game = GameState()

# 游戏主循环
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        
        # 键盘控制
        if event.type == pygame.KEYDOWN:
            if not game.game_over:
                if event.key == pygame.K_UP and game.direction != (0, 1):
                    game.next_direction = (0, -1)
                elif event.key == pygame.K_DOWN and game.direction != (0, -1):
                    game.next_direction = (0, 1)
                elif event.key == pygame.K_LEFT and game.direction != (1, 0):
                    game.next_direction = (-1, 0)
                elif event.key == pygame.K_RIGHT and game.direction != (-1, 0):
                    game.next_direction = (1, 0)
            else:
                if event.key == pygame.K_SPACE:
                    game.reset()
    
    # 更新游戏状态
    game.update()
    
    # 绘制游戏
    screen.fill(GRID_COLOR)
    
    # 绘制网格线
    for x in range(0, WINDOW_SIZE, CELL_SIZE):
        pygame.draw.line(screen, (200, 200, 200), (x, 0), (x, WINDOW_SIZE))
    for y in range(0, WINDOW_SIZE, CELL_SIZE):
        pygame.draw.line(screen, (200, 200, 200), (0, y), (WINDOW_SIZE, y))
    
    # 绘制食物
    if game.food:
        food_x, food_y = game.food
        pygame.draw.rect(screen, FOOD_COLOR, 
                         (food_x * CELL_SIZE, food_y * CELL_SIZE, CELL_SIZE, CELL_SIZE))
    
    # 绘制蛇
    for i, (x, y) in enumerate(game.snake):
        color = SNAKE_HEAD_COLOR if i == 0 else SNAKE_COLOR
        pygame.draw.rect(screen, color, 
                         (x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE))
    
    # 显示分数
    font = pygame.font.SysFont("华文宋体", 36)
    score_text = font.render(f"得分: {game.score}", True, TEXT_COLOR)
    screen.blit(score_text, (10, 10))
    
    # 游戏结束显示
    if game.game_over:
        font = pygame.font.SysFont("华文宋体", 72)
        game_over_text = font.render("游戏结束!", True, TEXT_COLOR)
        final_score_text = font.render(f"最终得分: {game.score}", True, TEXT_COLOR)
        restart_text = font.render("按空格键重新开始", True, TEXT_COLOR)
        
        screen.blit(game_over_text, 
                   (WINDOW_SIZE//2 - game_over_text.get_width()//2, WINDOW_SIZE//2 - 100))
        screen.blit(final_score_text, 
                   (WINDOW_SIZE//2 - final_score_text.get_width()//2, WINDOW_SIZE//2))
        screen.blit(restart_text, 
                   (WINDOW_SIZE//2 - restart_text.get_width()//2, WINDOW_SIZE//2 + 100))
    
    pygame.display.flip()
    clock.tick(10)  # 控制游戏速度